# -*- coding : utf-8 -*-

from docx import Document
from docx.shared import Inches
import docx2pdf as dp
import os,chardet

#------判断文本编码 begin------#
def detect_encoding(text):
    result = chardet.detect(text)
    encoding = result['encoding']
    confidence = result['confidence']
    return encoding, confidence
def getTxtCharset(filename:str):
    with open(filename, 'rb') as f:
        text = f.read()
    encoding, confidence = detect_encoding(text)
    return encoding
    #print(f"Detected encoding: {encoding}")
    #print(f"Confidence: {confidence}")
#------判断文本编码 end ------#

#读取srt文件
def getSublinesFromSrt(srtFile):
    lines = []
    _charset = getTxtCharset(srtFile)
    #print(_charset)
    with open(srtFile,'r',encoding=_charset) as f:
        lines = f.readlines()
    return lines

def deleteIllMark(txt:str):
    return txt.replace('\n','')

#解析srt文件list，组装成元组(序号，时间，英文，中文)
def processSublines(sublines):
    sublines_sum = len(sublines)
    targetList = []
    for i in range(0,sublines_sum):
        if '-->' in sublines[i]:
            orderNo = deleteIllMark(sublines[i-1])
            timeNo = deleteIllMark(sublines[i])
            englishSentence = deleteIllMark(sublines[i+2])
            chineseSentence = deleteIllMark(sublines[i+1])
            targetList.append((orderNo,timeNo,englishSentence,chineseSentence))
    return targetList

def generateDocWord(docx_dir:str,filename:str,titlename:str,subTupleList:list):
    document = Document()
    
    document.add_heading(titlename,0)
    
    records = subTupleList
    
    table = document.add_table(rows=1, cols=3,style='Medium Grid 3 Accent 4')
    hdr_cells = table.rows[0].cells
    hdr_cells[0].text = 'time'
    hdr_cells[1].text = 'English'
    hdr_cells[2].text = 'Chinese'
    hdr_cells[0].width = Inches(2)
    hdr_cells[1].width = Inches(5)
    hdr_cells[2].width = Inches(5)
    for orderNo, timeNo, english,chinese in records:
        row_cells = table.add_row().cells
        row_cells[0].text = getTimeNo(timeNo)
        row_cells[1].text = english
        row_cells[2].text = chinese
    
    #document.add_page_break()
    #添加页码
    add_page_number_by_p(document.sections[0].footer.paragraphs[0])
    document.save(os.path.join(docx_dir,filename+'.docx'))

#------字段显示内容 begin------#
#针对不同格式需求随时更改

#时间戳
def getTimeNo(time_no:str):
    return time_no.split('-->')[0].strip()
#文件名
def getFileName(file_name:str):
    if len(file_name.split('.'))<6:
        return file_name
    return '.'.join(file_name.split('.')[0:-6])
#标题名
def getTitleName(title_name:str):
    if len(title_name.split('.'))<6:
        return title_name
    return '.'.join(title_name.split('.')[0:-6])
#------字段显示内容 end------#

#srt转word文档
def processSrtToWordFiles(dirPath:str,targetPath:str):
    validDir(targetPath)
    for f in os.listdir(dirPath):
        print(f)
        srtLines = getSublinesFromSrt(dirPath+f)
        srtTupleLines = processSublines(srtLines)
        generateDocWord(targetPath,getFileName(f),getTitleName(f),srtTupleLines)
        #docx转pdf 若生成docx文件后不需要同时生成pdf就注释掉
        dp.convert(os.path.join(targetPath,getFileName(f)+'.docx'),validDir(os.path.join(targetPath,'pdf'))+'/'+getFileName(f)+'.pdf')
        #break #调试代码用

#批量docx转pdf
def processDocxToPdf(docx_dir:str,pdf_dir:str):
    tar_dir = validDir(pdf_dir)
    for d in os.listdir(docx_dir):
        print(d)
        dp.convert(os.path.join(docx_dir,d),os.path.join(tar_dir,d.replace('.docx','.pdf')))
        #break #调试代码用 正式运行时注释掉

#------添加页码 begin------#
from docx.oxml import OxmlElement, ns
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT

def create_element(name):
    return OxmlElement(name)

def create_attribute(element, name, value):
    element.set(ns.qn(name), value)


def add_page_number(run):
    fldChar1 = create_element('w:fldChar')
    create_attribute(fldChar1, 'w:fldCharType', 'begin')

    instrText = create_element('w:instrText')
    create_attribute(instrText, 'xml:space', 'preserve')
    instrText.text = "PAGE"

    fldChar2 = create_element('w:fldChar')
    create_attribute(fldChar2, 'w:fldCharType', 'end')

    run._r.append(fldChar1)
    run._r.append(instrText)
    run._r.append(fldChar2)

def add_page_number_by_p(paragraph):
    paragraph.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER

    page_run = paragraph.add_run()
    t1 = create_element('w:t')
    create_attribute(t1, 'xml:space', 'preserve')
    t1.text = '免费文档 禁止商用 Page '
    page_run._r.append(t1)

    page_num_run = paragraph.add_run()

    fldChar1 = create_element('w:fldChar')
    create_attribute(fldChar1, 'w:fldCharType', 'begin')

    instrText = create_element('w:instrText')
    create_attribute(instrText, 'xml:space', 'preserve')
    instrText.text = "PAGE"

    fldChar2 = create_element('w:fldChar')
    create_attribute(fldChar2, 'w:fldCharType', 'end')

    page_num_run._r.append(fldChar1)
    page_num_run._r.append(instrText)
    page_num_run._r.append(fldChar2)

    of_run = paragraph.add_run()
    t2 = create_element('w:t')
    create_attribute(t2, 'xml:space', 'preserve')
    t2.text = ' of '
    of_run._r.append(t2)

    fldChar3 = create_element('w:fldChar')
    create_attribute(fldChar3, 'w:fldCharType', 'begin')

    instrText2 = create_element('w:instrText')
    create_attribute(instrText2, 'xml:space', 'preserve')
    instrText2.text = "NUMPAGES"

    fldChar4 = create_element('w:fldChar')
    create_attribute(fldChar4, 'w:fldCharType', 'end')

    num_pages_run = paragraph.add_run()
    num_pages_run._r.append(fldChar3)
    num_pages_run._r.append(instrText2)
    num_pages_run._r.append(fldChar4)

#------添加页码 end  ------#

#若目录不存在 创建
def validDir(_dir:str):
    if not os.path.exists(_dir):
        os.makedirs(_dir)
    return _dir


if __name__ == '__main__' :
    SRT_DIR = 'C:/E/temp/pic/srt/'
    TAR_DIR = 'C:/E/temp/pic/'
    print('*'*10,'process start','*'*10)
    #lines = processSublines(getSublinesFromSrt('community.srt'))
    #generateDocWord('demo','biology',lines)
    #getSrtFiles(SRT_DIR)
    #print(validDir(os.path.join(TAR_DIR,'lrc','src')))
    processSrtToWordFiles(SRT_DIR,validDir(os.path.join(TAR_DIR,'words.')))
    #processDocxToPdf(os.path.join(TAR_DIR,'words'),os.path.join(TAR_DIR,'pdfs'))
    print('*'*10,'process end','*'*10)
